﻿// 确保在包含windows.h之前定义WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN  // 排除极少使用的内容
#include <winsock2.h>
#include <windows.h>
#include <ws2tcpip.h>
#include <iostream>
#include <vector>
#include <thread>
#include <atomic>
#include <cstring>
#include "Constants.h"

#pragma comment(lib, "ws2_32.lib")

// 禁用 Winsock 1 相关的警告
#define _WINSOCK_DEPRECATED_NO_WARNINGS
// 防止 Windows.h 包含 min/max 宏
#define NOMINMAX

// 其他头文件和常量定义...

// 消息类型（与客户端保持一致）
enum MessageType {
    SREEN_INFO = 1,
    SREEN_DATA,
    MOU_EVENT,
    BOARD_EVENT
};

// 屏幕信息结构
struct ScreenInfo {
    int width;
    int height;
    int bpp;  // 位深度
};

// 鼠标事件结构
struct MouseEvent {
    int x;
    int y;
    int eventType;  // 0:移动, 1:左键按下, 2:左键释放
};

// 键盘事件结构
struct KeyboardEvent {
    unsigned int keyCode;
    bool isDown;
};

// 全局变量
std::atomic<bool> running(true);
SOCKET listenSocket = INVALID_SOCKET;
SOCKET clientSocket = INVALID_SOCKET;
HDC hdcScreen = NULL;
HBITMAP hBitmap = NULL;
HDC hdcMem = NULL;
BITMAPINFOHEADER bi = { 0 };
int screenWidth = 0;
int screenHeight = 0;
int screenBpp = 32;  // 默认32位

// 函数声明
void InitializeScreenCapture();
void CaptureAndSendScreen();
void HandleClientEvents();

int main() {
    WSADATA wsaData;
    int iResult;

    // 初始化 Winsock
    iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != 0) {
        std::cerr << "WSAStartup failed: " << iResult << std::endl;
        return 1;
    }

    // 创建监听套接字
    listenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (listenSocket == INVALID_SOCKET) {
        std::cerr << "socket failed: " << WSAGetLastError() << std::endl;
        WSACleanup();
        return 1;
    }

    // 绑定套接字
    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_addr.s_addr = INADDR_ANY;
    serverAddr.sin_port = htons(DEFAULT_PORT);

    iResult = bind(listenSocket, (sockaddr*)&serverAddr, sizeof(serverAddr));
    if (iResult == SOCKET_ERROR) {
        std::cerr << "bind failed: " << WSAGetLastError() << std::endl;
        closesocket(listenSocket);
        WSACleanup();
        return 1;
    }

    // 监听连接
    iResult = listen(listenSocket, SOMAXCONN);
    if (iResult == SOCKET_ERROR) {
        std::cerr << "listen failed: " << WSAGetLastError() << std::endl;
        closesocket(listenSocket);
        WSACleanup();
        return 1;
    }

    std::cout << "Server started. Waiting for connections..." << std::endl;

    // 接受客户端连接
    sockaddr_in clientAddr;
    int clientAddrLen = sizeof(clientAddr);
    clientSocket = accept(listenSocket, (sockaddr*)&clientAddr, &clientAddrLen);
    if (clientSocket == INVALID_SOCKET) {
        std::cerr << "accept failed: " << WSAGetLastError() << std::endl;
        closesocket(listenSocket);
        WSACleanup();
        return 1;
    }

    std::cout << "Client connected." << std::endl;

    // 初始化屏幕捕获
    InitializeScreenCapture();

    // 发送屏幕信息
    ScreenInfo screenInfo = { screenWidth, screenHeight, screenBpp };
    int messageType = SREEN_INFO;
    send(clientSocket, (char*)&messageType, sizeof(messageType), 0);
    send(clientSocket, (char*)&screenInfo, sizeof(screenInfo), 0);

    // 创建线程处理客户端事件
    std::thread eventThread(HandleClientEvents);

    // 主循环：捕获并发送屏幕数据
    while (running) {
        CaptureAndSendScreen();
        Sleep(30);  // 约30ms/帧，约33FPS
    }

    // 等待线程结束
    if (eventThread.joinable()) eventThread.join();

    // 清理资源
    if (hBitmap) DeleteObject(hBitmap);
    if (hdcMem) DeleteDC(hdcMem);
    if (hdcScreen) ReleaseDC(NULL, hdcScreen);
    if (clientSocket != INVALID_SOCKET) closesocket(clientSocket);
    if (listenSocket != INVALID_SOCKET) closesocket(listenSocket);
    WSACleanup();

    return 0;
}

// 初始化屏幕捕获
void InitializeScreenCapture() {
    // 获取屏幕尺寸
    screenWidth = GetSystemMetrics(SM_CXSCREEN);
    screenHeight = GetSystemMetrics(SM_CYSCREEN);

    // 获取屏幕DC
    hdcScreen = GetDC(NULL);
    hdcMem = CreateCompatibleDC(hdcScreen);

    // 创建位图
    hBitmap = CreateCompatibleBitmap(hdcScreen, screenWidth, screenHeight);
    SelectObject(hdcMem, hBitmap);

    // 初始化位图信息
    bi.biSize = sizeof(BITMAPINFOHEADER);
    bi.biWidth = screenWidth;
    bi.biHeight = screenHeight;  // 正数表示从下到上
    bi.biPlanes = 1;
    bi.biBitCount = screenBpp;
    bi.biCompression = BI_RGB;
}

// 捕获并发送屏幕数据
void CaptureAndSendScreen() {
    if (!running || clientSocket == INVALID_SOCKET) return;

    // 捕获屏幕
    BitBlt(hdcMem, 0, 0, screenWidth, screenHeight, hdcScreen, 0, 0, SRCCOPY);

    // 获取位图数据大小
    int dataSize = screenWidth * screenHeight * (screenBpp / 8);
    std::vector<char> screenData(dataSize);

    // 获取位图数据
    GetDIBits(hdcMem, hBitmap, 0, screenHeight, screenData.data(),
        (BITMAPINFO*)&bi, DIB_RGB_COLORS);

    // 发送屏幕数据
    int messageType = SREEN_DATA;
    send(clientSocket, (char*)&messageType, sizeof(messageType), 0);
    send(clientSocket, (char*)&dataSize, sizeof(dataSize), 0);
    send(clientSocket, screenData.data(), dataSize, 0);
}

// 处理客户端事件（鼠标/键盘）
void HandleClientEvents() {
    char buffer[BUFFER_SIZE];
    int bytesReceived;

    while (running && clientSocket != INVALID_SOCKET) {
        // 接收消息类型
        if ((bytesReceived = recv(clientSocket, buffer, sizeof(int), 0)) <= 0) {
            std::cerr << "Client disconnected." << std::endl;
            running = false;
            break;
        }

        int messageType = *(int*)buffer;

        switch (messageType) {
        case MOU_EVENT: {
            if ((bytesReceived = recv(clientSocket, buffer, sizeof(MouseEvent), 0)) <= 0) {
                running = false;
                break;
            }

            MouseEvent mouseEvent = *(MouseEvent*)buffer;

            // 模拟鼠标事件
            INPUT input = { 0 };
            input.type = INPUT_MOUSE;

            switch (mouseEvent.eventType) {
            case 0:  // 移动
                input.mi.dwFlags = MOUSEEVENTF_MOVE | MOUSEEVENTF_ABSOLUTE;
                input.mi.dx = (LONG)((mouseEvent.x * 65535) / screenWidth);
                input.mi.dy = (LONG)((mouseEvent.y * 65535) / screenHeight);
                SendInput(1, &input, sizeof(INPUT));
                break;

            case 1:  // 左键按下
                input.mi.dwFlags = MOUSEEVENTF_LEFTDOWN;
                SendInput(1, &input, sizeof(INPUT));
                break;

            case 2: {  // 左键释放
                input.mi.dwFlags = MOUSEEVENTF_LEFTUP;
                SendInput(1, &input, sizeof(INPUT));
                break;
            }
            }
            break;
        }

        case BOARD_EVENT:
            if ((bytesReceived = recv(clientSocket, buffer, sizeof(KeyboardEvent), 0)) <= 0) {
                running = false;
                break;
            }

            KeyboardEvent keyboardEvent = *(KeyboardEvent*)buffer;

            // 模拟键盘事件
            INPUT keyInput = { 0 };
            keyInput.type = INPUT_KEYBOARD;
            keyInput.ki.wVk = keyboardEvent.keyCode;
            keyInput.ki.dwFlags = keyboardEvent.isDown ? 0 : KEYEVENTF_KEYUP;
            SendInput(1, &keyInput, sizeof(INPUT));
            break;
        }
    }
}